package com.wxhntmy.google.oauth2;

import com.alibaba.fastjson2.JSONObject;
import com.wxhntmy.config.Oauth2;
import com.wxhntmy.config.ProjectConfig;
import com.wxhntmy.httpclient.RestMockByProxy;
import com.wxhntmy.redis.JedisLock;
import com.wxhntmy.redis.JedisTool;
import org.apache.commons.codec.binary.Base64;
import org.apache.commons.lang3.StringUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import redis.clients.jedis.JedisPooled;

import java.io.FileInputStream;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.util.HashMap;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

/**
 * oauth2 工具类
 *
 * @author chenwc
 * @date 2024/7/30
 */
public class GoogleTranslateOauth2Util {

    private static final Logger log = LoggerFactory.getLogger(GoogleTranslateOauth2Util.class);

    private static final Oauth2 OAUTH_2 = ProjectConfig.getInstance().getGoogle().getOauth2();

    /**
     * 加载私钥
     *
     * @param keystorePath 私钥文件路径
     * @param keystorePass 私钥文件密码
     * @param alias        私钥别名
     * @return 私钥
     * @throws Exception 异常
     */
    public static PrivateKey loadPrivateKeyFromJks(String keystorePath, String keystorePass, String alias) throws Exception {
        KeyStore keyStore = KeyStore.getInstance("JKS");
        keyStore.load(new FileInputStream(keystorePath), keystorePass.toCharArray());
        return (PrivateKey) keyStore.getKey(alias, keystorePass.toCharArray());
    }

    /**
     * 获取当前时间戳
     *
     * @return 时间戳
     */
    public static long getTimeStamp() {

        return System.currentTimeMillis() / 1000;
    }

    /**
     * 生成JWT
     *
     * @param privateKey 私钥
     * @return JWT
     */
    public static String generateJwt(PrivateKey privateKey) {

        JSONObject headers = new JSONObject();
        headers.put("alg", "RS256");
        headers.put("typ", "JWT");
        headers.put("kid", OAUTH_2.getPrivateKeyId());
        log.debug("headers: " + headers.toJSONString());
        String base64Header = Base64.encodeBase64String(headers.toJSONString().getBytes(StandardCharsets.UTF_8));
        log.debug("base64Header: " + base64Header);
        JSONObject payload = new JSONObject();
        payload.put("iss", OAUTH_2.getClientId());
        payload.put("scope", String.join(" ", OAUTH_2.getScopeUrl()));
        payload.put("aud", OAUTH_2.getAudUrl());
        long timeStamp = getTimeStamp();
        payload.put("iat", timeStamp);
        // 缓存过期时间
        long expireDate = timeStamp + 3600;
        setExpireDateToRedis(expireDate);
        payload.put("exp", expireDate);
        log.debug("payload: " + payload.toJSONString());
        String base64Payload = Base64.encodeBase64String(payload.toJSONString().getBytes(StandardCharsets.UTF_8));
        log.debug("base64Payload: " + base64Payload);


        // 创建Signature实例并初始化
        Signature signature;
        String signatureStr;
        try {
            signature = Signature.getInstance("SHA256withRSA");
            signature.initSign(privateKey);
            // 需要签名的数据
            String dataStr = base64Header + "." + base64Payload;
            byte[] data = dataStr.getBytes(StandardCharsets.UTF_8);
            // 更新签名对象的数据
            signature.update(data);
            // 计算签名
            byte[] signedData = signature.sign();
            signatureStr = Base64.encodeBase64String(signedData);
            // 打印签名结果，通常会将它转换成Base64编码以便传输
            log.debug("signature: " + signatureStr);
        } catch (NoSuchAlgorithmException | InvalidKeyException | SignatureException e) {
            throw new RuntimeException(e);
        }
        String jwt = base64Header + "." + base64Payload + "." + signatureStr;
        log.debug("jwt: " + jwt);

        return jwt;
    }

    /**
     * 获取accessToken
     *
     * @return accessToken
     */
    public static String getAccessTokenFromGoogle() {
        PrivateKey privateKey;
        try {
            privateKey = loadPrivateKeyFromJks(OAUTH_2.getJksStorePath(), OAUTH_2.getJksPassword(), OAUTH_2.getKeyAlias());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
        String jwt = generateJwt(privateKey);
        Map<String, String> form = new HashMap<>();
        form.put("grant_type", "urn:ietf:params:oauth:grant-type:jwt-bearer");
        form.put("assertion", jwt);
        Map<String, String> header = new HashMap<>();
        header.put("Content-Type", "application/x-www-form-urlencoded");
        String rst = RestMockByProxy.sendPostByFormBody(OAUTH_2.getTokenUrl(), header, form);
        log.info("生成JWT结果: " + rst);
        return JSONObject.parseObject(rst).getString("access_token");
    }

    /**
     * 判断accessToken是否过期
     *
     * @return 是否过期
     */
    public static boolean isExpireFromRedis() {
        String expireDateStr;
        JedisPooled jedis = JedisTool.getJedis();
        expireDateStr = jedis.get("expireDate");
        //JedisTool.close();
        if (StringUtils.isEmpty(expireDateStr)) {
            expireDateStr = "0";
        }
        long expireDate = Long.parseLong(expireDateStr);
        log.info("expireDate: {}", expireDate);
        return getTimeStamp() > expireDate;
    }

    /**
     * 设置过期时间
     *
     * @param expireDate 过期时间
     */
    public static void setExpireDateToRedis(long expireDate) {
        JedisPooled jedis = JedisTool.getJedis();
        JedisLock demoEtcdLock1 = new JedisLock(jedis, "expireDate-lock", UUID.randomUUID().toString(), 3599);
        boolean lock1 = demoEtcdLock1.tryLock(20, TimeUnit.SECONDS);
        if (lock1) {
            try {
                jedis.set("expireDate", String.valueOf(expireDate));
                jedis.expire("expireDate", 3599);
                log.debug("get lock1");
            } finally {
                demoEtcdLock1.unlock();
            }
        }
        //JedisTool.close();
    }

    /**
     * 获取accessToken
     *
     * @return accessToken
     */
    public static String getAccessToken() {
        String accessToken;
        if (isExpireFromRedis()) {
            accessToken = getAccessTokenFromGoogle();
            setAccessTokenToRedis(accessToken);
            log.debug("accessToken: {}", accessToken);
        } else {
            accessToken = getAccessTokenFromRedis();
        }
        return accessToken;
    }

    /**
     * 设置accessToken
     *
     * @param accessToken accessToken
     */
    public static void setAccessTokenToRedis(String accessToken) {
        JedisPooled jedis = JedisTool.getJedis();
        JedisLock demoEtcdLock1 = new JedisLock(jedis, "accessToken-lock", UUID.randomUUID().toString(), 15);
        boolean lock1 = demoEtcdLock1.tryLock(20, TimeUnit.SECONDS);
        if (lock1) {
            try {
                jedis.set("accessToken", accessToken);
                jedis.expire("accessToken", 3599);
            } finally {
                demoEtcdLock1.unlock();
            }
        }
        //JedisTool.close();
    }

    /**
     * 获取accessToken
     *
     * @return accessToken
     */
    public static String getAccessTokenFromRedis() {
        String accessToken;
        JedisPooled jedis = JedisTool.getJedis();
        accessToken = jedis.get("accessToken");
        if (StringUtils.isEmpty(accessToken)) {
            accessToken = getAccessTokenFromGoogle();
            setAccessTokenToRedis(accessToken);
            log.debug("accessToken: {}", accessToken);
        }
        //JedisTool.close();
        return accessToken;
    }
}
